#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
模板匹配应用程序 - 模板匹配视图模型
实现OpenCV模板匹配功能的核心类
"""

import os
import cv2
import numpy as np
from PyQt5.QtGui import QImage, QPixmap

class TemplateMatchingViewModel:
    """
    模板匹配视图模型
    处理图像加载、模板匹配和结果生成
    """
    
    def __init__(self):
        """初始化视图模型"""
        # 图像属性
        self.source_image = None  # 原始图像
        self.template_image = None  # 模板图像
        self.result_image = None  # 结果图像
        self.heatmap_image = None  # 热图结果
        
        # OpenCV格式图像
        self.source_mat = None  # 原始图像的Mat对象
        self.template_mat = None  # 模板图像的Mat对象
        
        # 匹配结果数据
        self.last_match_results = {}
        
        
        # 匹配方法列表
        self.matching_methods = [
            ("平方差匹配法", cv2.TM_SQDIFF),
            ("归一化平方差匹配法", cv2.TM_SQDIFF_NORMED),
            ("相关匹配法", cv2.TM_CCORR),
            ("归一化相关匹配法", cv2.TM_CCORR_NORMED),
            ("相关系数匹配法", cv2.TM_CCOEFF),
            ("归一化相关系数匹配法", cv2.TM_CCOEFF_NORMED)
        ]
        
        # 默认使用归一化相关系数匹配法
        self.current_method_index = 5
        
    def load_source_image(self, image_path):
        """
        加载源图像
        
        参数:
            image_path: 图像文件路径
        """
        try:
            # 使用OpenCV加载图像
            self.source_mat = cv2.imread(image_path)
            if self.source_mat is None:
                return None
                
            # 转换为RGB格式（OpenCV加载为BGR）
            rgb_image = cv2.cvtColor(self.source_mat, cv2.COLOR_BGR2RGB)
            
            # 创建QImage
            height, width, channels = rgb_image.shape
            bytes_per_line = channels * width
            q_image = QImage(rgb_image.data, width, height, bytes_per_line, QImage.Format_RGB888)
            
            # 保存为QPixmap
            self.source_image = QPixmap.fromImage(q_image)
            return self.source_image
        except Exception as e:
            print(f"加载图像失败: {str(e)}")
            return None
            
    def load_template_image(self, image_path):
        """
        加载模板图像
        
        参数:
            image_path: 模板图像文件路径
        """
        try:
            # 使用OpenCV加载图像
            self.template_mat = cv2.imread(image_path)
            if self.template_mat is None:
                return None
                
            # 转换为RGB格式（OpenCV加载为BGR）
            rgb_image = cv2.cvtColor(self.template_mat, cv2.COLOR_BGR2RGB)
            
            # 创建QImage
            height, width, channels = rgb_image.shape
            bytes_per_line = channels * width
            q_image = QImage(rgb_image.data, width, height, bytes_per_line, QImage.Format_RGB888)
            
            # 保存为QPixmap
            self.template_image = QPixmap.fromImage(q_image)
            return self.template_image
        except Exception as e:
            print(f"加载模板图像失败: {str(e)}")
            return None
            
    def set_template_image(self, template_pixmap):
        """
        设置模板图像 (从QPixmap)
        
        参数:
            template_pixmap: QPixmap格式的模板图像
        """
        if template_pixmap is None:
            return
            
        # 保存QPixmap
        self.template_image = template_pixmap
        
        # 转换为OpenCV格式
        qimage = template_pixmap.toImage()
        width = qimage.width()
        height = qimage.height()
        
        # 将QImage转换为numpy数组
        ptr = qimage.bits()
        ptr.setsize(qimage.byteCount())
        arr = np.array(ptr).reshape(height, width, 4)  # RGBA
        
        # 转换为BGR (OpenCV格式)
        self.template_mat = cv2.cvtColor(arr, cv2.COLOR_RGBA2BGR)
        
    def save_template_image(self, file_path):
        """
        保存模板图像
        
        参数:
            file_path: 保存路径
        """
        if self.template_image is None:
            return False
            
        return self.template_image.save(file_path)
        
    def match_template(self, method_index=None):
        """
        执行模板匹配
        
        参数:
            method_index: 匹配方法索引，None时使用当前方法
        """
        if self.source_mat is None or self.template_mat is None:
            return None
            
        # 使用指定方法或当前方法
        if method_index is not None:
            self.current_method_index = method_index
            
        # 获取当前匹配方法
        method_name, method = self.matching_methods[self.current_method_index]
        
        # 复制原图，用于绘制结果
        display_image = self.source_mat.copy()
        
        # 转为灰度图进行处理
        gray_source = cv2.cvtColor(self.source_mat, cv2.COLOR_BGR2GRAY)
        gray_template = cv2.cvtColor(self.template_mat, cv2.COLOR_BGR2GRAY)
        
        # 执行模板匹配
        result = cv2.matchTemplate(gray_source, gray_template, method)
        
        # 查找最佳匹配位置
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
        
        # 确定最佳匹配位置和分数
        # 对于TM_SQDIFF或TM_SQDIFF_NORMED方法，最小值为最佳匹配
        if method == cv2.TM_SQDIFF or method == cv2.TM_SQDIFF_NORMED:
            match_loc = min_loc
            score = min_val
        else:
            match_loc = max_loc
            score = max_val
            
        # 计算中心点
        h, w = self.template_mat.shape[:2]
        center_x = match_loc[0] + w // 2
        center_y = match_loc[1] + h // 2
        
        # 绘制匹配结果矩形
        cv2.rectangle(
            display_image,
            match_loc,
            (match_loc[0] + w, match_loc[1] + h),
            (0, 255, 0),
            2
        )
        
        # 绘制中心十字线
        self._draw_crosshair(display_image, (center_x, center_y), (0, 255, 0), 20, 2)
        
        # 计算结果的统计信息
        mean, std_dev = cv2.meanStdDev(result)
        
        # 存储匹配结果数据
        self.last_match_results = {
            'method': method,
            'method_name': method_name,
            'score': float(score),
            'position': (center_x, center_y),
            'rect': (match_loc[0], match_loc[1], match_loc[0] + w, match_loc[1] + h),
            'mean': float(mean[0][0]),
            'std_dev': float(std_dev[0][0]),
            'template_size': (w, h)
        }
        
        # 在图像上显示信息
        info_text = [
            f"Method: {method_name}",
            f"Score: {score:.4f}",
            f"Position: ({center_x}, {center_y})",
            f"Size: {w}x{h}",
            f"Mean: {mean[0][0]:.4f}",
            f"StdDev: {std_dev[0][0]:.4f}"
        ]
        
        # 绘制信息文本
        for i, text in enumerate(info_text):
            cv2.putText(
                display_image,
                text,
                (10, 30 + i * 25),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.7,
                (0, 0, 255),
                1,
                cv2.LINE_AA
            )
            
        # 生成热图
        self._generate_heatmap(result, score, method)
        
        # 简化代码，移除多匹配点功能，增强稳定性
        
        # 转换结果为QPixmap
        rgb_result = cv2.cvtColor(display_image, cv2.COLOR_BGR2RGB)
        height, width, channels = rgb_result.shape
        bytes_per_line = channels * width
        q_image = QImage(rgb_result.data, width, height, bytes_per_line, QImage.Format_RGB888)
        self.result_image = QPixmap.fromImage(q_image)
        
        return self.result_image
        
    def _generate_heatmap(self, result, score, method):
        """
        生成匹配结果热图
        
        参数:
            result: 匹配结果矩阵
            score: 最佳匹配分数
            method: 使用的匹配方法
        """
        # 标准化结果到[0, 255]范围
        result_normalized = cv2.normalize(result, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
        
        # 对于TM_SQDIFF方法，需要反转结果
        if method == cv2.TM_SQDIFF or method == cv2.TM_SQDIFF_NORMED:
            result_normalized = 255 - result_normalized
            
        # 应用色彩映射
        heatmap = cv2.applyColorMap(result_normalized, cv2.COLORMAP_JET)
        
        # 绘制坐标网格（每50个像素一条线）
        h, w = heatmap.shape[:2]
        grid_spacing = 50
        for x in range(0, w, grid_spacing):
            cv2.line(heatmap, (x, 0), (x, h-1), (255, 255, 255), 1)
        for y in range(0, h, grid_spacing):
            cv2.line(heatmap, (0, y), (w-1, y), (255, 255, 255), 1)
        
        # 添加分数信息
        cv2.putText(
            heatmap,
            f"Best Score: {score:.4f}",
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (255, 255, 255),
            1,
            cv2.LINE_AA
        )
        
        # 转换为QPixmap
        rgb_heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
        height, width, channels = rgb_heatmap.shape
        bytes_per_line = channels * width
        q_image = QImage(rgb_heatmap.data, width, height, bytes_per_line, QImage.Format_RGB888)
        self.heatmap_image = QPixmap.fromImage(q_image)
        
    # 移除多匹配点功能，简化代码
    
    def _draw_crosshair(self, image, center, color, size, thickness):
        """
        绘制十字准星
        
        参数:
            image: 要绘制的图像
            center: 中心点坐标 (x, y)
            color: 颜色 (B, G, R)
            size: 十字线的长度
            thickness: 线宽
        """
        # 绘制十字准星圈
        cv2.circle(image, center, size // 2, color, 1)
        
        # 绘制水平线
        cv2.line(
            image,
            (center[0] - size, center[1]),
            (center[0] + size, center[1]),
            color,
            thickness
        )
        
        # 绘制垂直线
        cv2.line(
            image,
            (center[0], center[1] - size),
            (center[0], center[1] + size),
            color,
            thickness
        )
